package com.example.springbootdemo.text.zh;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.wltea.analyzer.lucene.IKAnalyzer;
import java.io.IOException;
import java.io.StringReader;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

public class WordsHelper {

    private Analyzer analyzer = new IKAnalyzer(true);

	/**
	 * 分词
	 * @param str
	 * @return
	 */
	public Set<String> parseWords(String str) {
		return this.parseWordCounts(str).keySet();
	}

	/**
	 * 分词
	 * @param str
	 * @return
	 */
	public Map<String, Integer> parseWordCounts(String str) {

		// 获取Lucene的TokenStream对象
		TokenStream ts = null;
		try {
			// 构建IK分词器，使用smart分词模式
			ts = analyzer.tokenStream("myfield", new StringReader(str));
			// 获取词元文本属性
			CharTermAttribute term = ts.addAttribute(CharTermAttribute.class);

			// 重置TokenStream（重置StringReader）
			ts.reset();
			Map<String, Integer> map = new HashMap<>();
			// 迭代获取分词结果
			while (ts.incrementToken()) {
				String word = term.toString();
				Integer i = map.get(word);
				if (i == null) {
					i = 0;
				}
				i = i + 1;
				map.put(word, i);
			}
			// 关闭TokenStream（关闭StringReader）
			ts.end(); // Perform end-of-stream operations, e.g. set the final offset.
			return map;
		} catch (IOException e) {
			throw new RuntimeException(e);
		} finally {
			// 释放TokenStream的所有资源
			if (ts != null) {
				try {
					ts.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
	}
}
